// com/tidu/strategy/task/ParallelWriteManager.java
package com.tidu.strategy.task;

import com.tidu.utils.DbUtilsTemplate;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ParallelWriteManager {
    private final ExecutorService executorService;
    private final DbUtilsTemplate targetDbUtils;
    private final String tableName;
    private final int threadCount;
    private final int batchSizePerThread;

    public ParallelWriteManager(DbUtilsTemplate targetDbUtils, String tableName, int threadCount, int batchSizePerThread) {
        this.targetDbUtils = targetDbUtils;
        this.tableName = tableName;
        this.threadCount = threadCount;
        this.batchSizePerThread = batchSizePerThread;
        this.executorService = Executors.newFixedThreadPool(threadCount);
    }

    public long writeDataInParallel(List<Map<String, Object>> dataList) throws Exception {
        if (dataList.isEmpty()) {
            return 0;
        }

        // 获取列信息
        List<String> columns = new ArrayList<>(dataList.get(0).keySet());

        // 将大数据集拆分成多个小批次
        List<List<Map<String, Object>>> batches = splitDataIntoBatches(dataList, batchSizePerThread);

        log.debug("Split {} records into {} batches for parallel processing", dataList.size(), batches.size());

        // 创建并行任务
        List<Future<Integer>> futures = new ArrayList<>();

        for (List<Map<String, Object>> batch : batches) {
            ParallelWriteTask task = new ParallelWriteTask(batch, targetDbUtils, tableName, columns);
            futures.add(executorService.submit(task));
        }

        // 等待所有任务完成并收集结果
        long totalWritten = 0;
        for (Future<Integer> future : futures) {
            totalWritten += future.get();
        }

        log.debug("Parallel write completed. Total records written: {}", totalWritten);
        return totalWritten;
    }

    private List<List<Map<String, Object>>> splitDataIntoBatches(List<Map<String, Object>> dataList, int batchSize) {
        List<List<Map<String, Object>>> batches = new ArrayList<>();

        for (int i = 0; i < dataList.size(); i += batchSize) {
            int end = Math.min(i + batchSize, dataList.size());
            batches.add(new ArrayList<>(dataList.subList(i, end)));
        }

        return batches;
    }

    public void shutdown() {
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
}
